package com.zhlq.util;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileWriter;
import java.io.IOException;
import java.io.Writer;
import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.SQLException;

import org.dbunit.DatabaseUnitException;
import org.dbunit.database.DatabaseConnection;
import org.dbunit.database.IDatabaseConnection;
import org.dbunit.database.QueryDataSet;
import org.dbunit.dataset.CachedDataSet;
import org.dbunit.dataset.IDataSet;
import org.dbunit.dataset.stream.IDataSetProducer;
import org.dbunit.dataset.xml.FlatXmlDataSet;
import org.dbunit.dataset.xml.XmlDataSetWriter;
import org.dbunit.dataset.xml.XmlProducer;
import org.dbunit.operation.DatabaseOperation;
import org.dbunit.util.FileHelper;

public final class DbUnitUtil {

	private DbUnitUtil() {
		super();
	}

	/**
	 * 导出数据到指定文件
	 * 
	 * @param file
	 *            一个标准的java.io.File
	 * @param connection
	 *            一个标准的java.sql.Connection
	 * @throws org.dbunit.DatabaseUnitException
	 */
	public static void export(File file, Connection connection, String table)
			throws DatabaseUnitException, IOException {
		IDatabaseConnection databaseConnection = new DatabaseConnection(connection);
		QueryDataSet dataSet = new QueryDataSet(databaseConnection);
		dataSet.addTable(table);
		Writer writer = new FileWriter(file);
		XmlDataSetWriter w = new XmlDataSetWriter(writer);
		w.write(dataSet);
		writer.flush();
		writer.close();
	}

	/**
	 * 导入数据到ROOM表
	 * 
	 * @param file
	 *            一个标准的java.io.File
	 * @param connection
	 *            一个标准的java.sql.Connection
	 */
	public static void imports(File file, Connection connection)
			throws DatabaseUnitException, IOException, SQLException {
		IDataSetProducer dataSetProducer = new XmlProducer(FileHelper.createInputSource(file));
		IDataSet dataSet = new CachedDataSet(dataSetProducer);
		IDatabaseConnection databaseConnection = new DatabaseConnection(connection);
		DatabaseOperation operation = DatabaseOperation.CLEAN_INSERT;
		DatabaseOperation.TRANSACTION(operation);
		operation.execute(databaseConnection, dataSet);
		DatabaseOperation.CLOSE_CONNECTION(operation);
	}
	
	public IDatabaseConnection getConnection()
			throws Exception {
		Class.forName("org.gjt.mm.mysql.Driver");
		Connection jdbcConnection = DriverManager.getConnection("jdbc:mysql://127.0.0.1/hr", "hr", "hr");
		return new DatabaseConnection(jdbcConnection);
	}
	
	@SuppressWarnings("deprecation")
	public IDataSet getDataSet() throws Exception {
		return new FlatXmlDataSet(new FileInputStream("hr-seed.xml"));
	}
	
	public DatabaseOperation getSetUpOperation()
			throws Exception {
		return DatabaseOperation.REFRESH;
	}
	
	public DatabaseOperation getTearDownOperation()
			throws Exception {
		return DatabaseOperation.NONE;
	}
	
}
